clear;
rng(50, 'twister')
if ispc; dbstop if error; end

if ~ismac && ~ispc 
    name = getenv('HOSTNAME');
    if contains(name, 'umich') % on UM cluster            
        np = str2num(getenv('SLURM_CPUS_PER_TASK')); 
        thePool = parpool('local', np);        
    else %Maryland
        parpool('local', maxNumCompThreads);
    end
end

%% monte carlo simulation setup
monte_carlo_exp = 1; % 1: monte carlo for coverage prob; 2: monte carlo for time test

if monte_carlo_exp ==1 
    N_vec = 2; % number of firms
else
    N_vec = 2:8; 
end
phi = 0.4; % competition effect parameter

N_monte = 300;
N_mkt = 4000; 

correlated_fc = false;

mktsize_lb = 1; mktsize_ub = 2;

% fixed cost parameters 
if correlated_fc
    para00 = [1 1 1];
else
    para00 = [1 1];
end

if monte_carlo_exp == 1 && ~isscalar(N_vec); error('when monte_carlo_exp = 1, N_vec needs to be a scalar'); end;

%% estimation set up
N_sim = 500;  % number of simulation draws for CT estimation
if correlated_fc; N_corr_sim = 100; end; 

tau_reps = 250;
sig_level = 0.05;
rand_seed = 100;

% moment choices
if monte_carlo_exp == 1
    N = N_vec;
    if N == 2
        profit_cutoff_min = (0.25:0.25:1.25);
        profit_cutoff_max = (0.25:0.25:1.25);
    elseif N == 3 || N == 4
        profit_cutoff_min = (0.25:0.25:1);
        profit_cutoff_max = (0.25:0.25:1);
    end
    if correlated_fc
        profit_cutoff_corr = 0.5;
    end
elseif monte_carlo_exp == 2
    profit_cutoff_min = (0.25:0.25:1.25);
    profit_cutoff_max = (0.25:0.25:1.25);
end

% grid points of parameter values
if monte_carlo_exp == 1
    Xgrid = para00(1) + (-2:0.125:2);
    Ygrid = log(para00(2)) + (-2:0.125:2); 
    if correlated_fc
        Zgrid = log(para00(3)) + (-2:0.125:2); 
        [Xmesh,Ymesh, Zmesh] = meshgrid(Xgrid, Ygrid, Zgrid);
    else
        Zgrid = []; Zmesh = [];
        [Xmesh,Ymesh] = meshgrid(Xgrid, Ygrid);
    end
    if correlated_fc
        para_grid = [Xmesh(:), exp(Ymesh(:)), exp(Zmesh(:))];
    else
        para_grid = [Xmesh(:), exp(Ymesh(:))];
    end
end
    
if monte_carlo_exp == 2    
    fy_t_as10 = nan(length(N_vec), N_monte);
    ct_t_as10 = nan(length(N_vec), N_monte);
    at_t_as10 = nan(length(N_vec), N_monte);
end

for nn = 1:length(N_vec)
    N = N_vec(nn);

    %% preparations
    % enumerate all possible product configuration: N_firm_outcome x N
    firm_outcome = [0; 1];
    for k = 1 : N-1
        firm_outcome = [[firm_outcome; firm_outcome], ...
            [zeros(size(firm_outcome, 1), 1); ones(size(firm_outcome, 1), 1)]];
    end
    N_firm_outcome = size(firm_outcome, 1);
    
    ind_unilateral_deivation = nan(size(firm_outcome)); % index for the outcome in firm_outcome that corresponds to a unilateral deviation: N_firm_outcome x N    
    for i = 1:N_firm_outcome 
        for k = 1:N
            tmp = firm_outcome(i, :);
            tmp(k) = 1 - tmp(k);
            ind_unilateral_deivation(i,k) = find(ismember(firm_outcome, tmp, 'rows')); % find the index of the outcome where firm k unilaterally deviates
        end
    end
    
    % draw covariates and fixed costs
    x1 = rand(N, N_mkt, N_monte); % own profit shiftor
    x2 = 0.5* rand(N, N_mkt, N_monte); % competitive effect
    mktsize = mktsize_lb + (mktsize_ub - mktsize_lb)*abs(rand(1, N_mkt, N_monte));
    if correlated_fc
        fc =  para00(1) + para00(2)*randn(N, N_mkt, N_monte) + para00(3)*randn(1, N_mkt, N_monte);
    else
        fc =  para00(1) + para00(2)*randn(N, N_mkt, N_monte);
    end

    % initialize estimation results
    if monte_carlo_exp == 1
        is_covered_ct_as10 = false(N_monte, 1); % whether the true parameter is covered, based on CT bounds and Andrews and Soares
        is_covered_at_as10 = false(N_monte, 1); % whether the true parameter is covered, based on AT bounds and Andrews and Soares
        is_covered_fy_as10 = false(N_monte, 1); % whether the true parameter is covered, based on FY bounds and Andrews and Soares
        if correlated_fc
            covered_mesh_ct_as10 = repmat(false(size(Xmesh)), [1 1 1 N_monte]); % coverage prob of a grid of points, AT bounds, Andrews and Soares
            covered_mesh_at_as10 = repmat(false(size(Xmesh)), [1 1 1 N_monte]); % coverage prob of a grid of points, CT bounds, Andrews and Soares
            covered_mesh_fy_as10 = repmat(false(size(Xmesh)), [1 1 1 N_monte]); % coverage prob of a grid of points, FY bounds, Andrews and Soares
        else
            covered_mesh_ct_as10 = repmat(false(size(Xmesh)), [1 1 N_monte]); % coverage prob of a grid of points, AT bounds, Andrews and Soares
            covered_mesh_at_as10 = repmat(false(size(Xmesh)), [1 1 N_monte]); % coverage prob of a grid of points, CT bounds, Andrews and Soares
            covered_mesh_fy_as10 = repmat(false(size(Xmesh)), [1 1 N_monte]); % coverage prob of a grid of points, FY bounds, Andrews and Soares
        end
    end

    for m = 1:N_monte  
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        %% Generate data
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        tic; 
        x1_m = x1(:, :, m);
        x2_m = x2(:, :, m);
        
        % compute each firm's variable profit and total profit in each outcome
        profit = nan(N_firm_outcome, N_mkt, N);
        totalprofit = nan(N_firm_outcome, N_mkt, N);
        for k = 1:N
            ind_not_k = [(1:k-1),(k+1:N)]; % index of competiting firms
            tmp = firm_outcome(:, ind_not_k) * x2_m(ind_not_k, :); % N_firm_outcome x N_mkt: sum across competing firms that are present in a market
            profit(:, :, k) = mktsize(1, :, m) .* (x1_m(k,:) - phi * log(1 + tmp)); 
    
            totalprofit(:, :, k) = firm_outcome(:, k) .* (profit(:, :, k) - fc(k, :, m));
        end
    
        % check whether each firm has an incentive to deviate in each outcome
        % check whether each firm is using a dominated strategy in each outcome
        exist_profit_deviation = false(N_firm_outcome, N_mkt, N);
        is_dominated = false(N_firm_outcome, N_mkt, N);
        for k = 1:N
             ind_not_k = [(1:k-1),(k+1:N)];
             for i = 1:N_firm_outcome
                tmp = firm_outcome(i, :);
                tmp(k) = 1 - tmp(k);
                ind = ismember(firm_outcome, tmp, 'rows'); % find the index of the outcome where firm k unilaterally deviates
                if nnz(ind) ~= 1; error('there should be a unique row in firm_outcome that corresponds to such an outcome'); end
                exist_profit_deviation(i, :, k) = totalprofit(i, :, k) < totalprofit(ind, :, k);
             end
           
            ind0 = firm_outcome(:, k) == 0;
            ind1 = ~ind0;
            if ~isequal(firm_outcome(ind0, ind_not_k), firm_outcome(ind1, ind_not_k))
                error('the following codes assume that firm_outcome(ind0, ind_not_k) = firm_outcome(ind1, ind_not_k)')
            end
        
            is_dominated(ind0, :, k) = repmat(all(0 < totalprofit(ind1, :, k), 1), [nnz(ind0), 1]); % action 0 is dominated
            is_dominated(ind1, :, k) = repmat(all(totalprofit(ind1, :, k) < 0, 1), [nnz(ind1), 1]); % action 1 is dominated
        end
        
        % check whether each outcome is a pure-strategy eqm and whether each outcome is a level-1 rational eqm
        is_equi = ~any(exist_profit_deviation, 3); % N_firm_outcome x N_mkt
        is_level1 = ~any(is_dominated, 3); % N_firm_outcome x N_mkt
            
        exist_profit_deviation = []; is_dominated = []; totalprofit = []; % clear some variables
    
        % initiate matrices that save the selected equilibrium
        data_outcome = nan(N_mkt, N); %simulated outcomes
        is_equi_data = false(N_firm_outcome, N_mkt); %whether the outcome in the first dim is the eqm in data_outcome
        
        n_equi = sum(is_equi, 1); % # of pure-strategy eqm: 1 x N_mkt
        
        % unique eqm
        ind_unique = find(n_equi == 1);
        for j = 1:length(ind_unique)
            jj = ind_unique(j);
            data_outcome(jj, :) = firm_outcome(is_equi(:, jj), :);
        end
        is_equi_data(:, ind_unique) = is_equi(:, ind_unique);
        
        % multiple eqm, select one randomly
        ind_multiple = find(n_equi > 1);
        for j = 1:length(ind_multiple)
            jj = ind_multiple(j);
            rand_equi_ind = randsample(find(is_equi(:, jj)), 1);
            data_outcome(jj, :) = firm_outcome(rand_equi_ind, :);
            is_equi_data(:, jj) =  false(N_firm_outcome, 1);
            is_equi_data(rand_equi_ind, jj) = true;
        end
            
        % no pure-strategy eqm, randomly select a level-1 rational strategy profile
        ind_no_pure = find(n_equi == 0);
        for j = 1:length(ind_no_pure)
            jj = ind_no_pure(j);
            rand_equi_ind = randsample(find(is_level1(:, jj)), 1);
            data_outcome(jj, :) = firm_outcome(rand_equi_ind, :);
            is_equi_data(:, jj) =  false(N_firm_outcome, 1);
            is_equi_data(rand_equi_ind, jj) = true;
        end
        
        is_equi = []; is_level1 = [];  % clear some variables
    
        fprintf(1, 'MC smln %1.0f, smln takes %1.2f sec\n', m, toc);
        
        % simulated outcome that will be used for estimation
        %   data_outcome: N_mkt x N
        %   profit: N_firm_outcome x N_mkt x N, variable profit
        %   is_equi_data: N_firm_outcome x N_mkt 
            
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % estimation
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        % min and max of the variable profit
        profit_lb = min(profit, [], 1); % 1 x N_mkt x N
        profit_lb = squeeze(profit_lb); % N_mkt x N
        profit_ub = max(profit, [], 1); % 1 x N_mkt x N
        profit_ub = squeeze(profit_ub); % N_mkt x N
    
        %% FY estimation
        profit_lb_vec = profit_lb(:);
        profit_ub_vec = profit_ub(:);
        data_outcome_vec = data_outcome(:);
    
        % instrument at the firm/market level
        g_ar_as10 = {gen_iv_monte_carlo(profit_lb_vec, profit_ub_vec, profit_cutoff_min, profit_cutoff_max)};
    
        if correlated_fc
            corr_script_monte_carlo
        else
            data_outcome_vec2 = [];    
            N_dev2 = [];
            dev2_ind_f1_tr = [];
            dev2_ind_f2_tr = [];
            corr_g_ar = [];
            corr_sim = []; 
            corr_sim_ct = [];
            corr_sim_at = [];        
        end
    
        % coverage rates & speed test   
        % (1) at the true parameter value    
        tic;
        [Tn00, cv00] = obj_fy_as10_corr(para00, profit_ub_vec, ... 
            profit_lb_vec, data_outcome_vec, g_ar_as10,  N_mkt, ...
            rand_seed, tau_reps, sig_level, ...
            correlated_fc, data_outcome_vec2, N_dev2, corr_sim, ...
            dev2_ind_f1_tr, dev2_ind_f2_tr, corr_g_ar);
    
        if monte_carlo_exp == 1
            is_covered_fy_as10(m) = Tn00<cv00;
        else
            fy_t_as10(nn, m) = toc;
        end
    
        % (2) at the parameter grid points
        if monte_carlo_exp==1 % monte carlo for coverage
            Tnmesh_fy_as10 = nan(size(Xmesh)); cvmesh_fy_as10 = nan(size(Xmesh));
    
            parfor n = 1 : size(para_grid, 1)
                [Tn, cv] = obj_fy_as10_corr(...
                    para_grid(n, :), profit_ub_vec, ...
                    profit_lb_vec, data_outcome_vec, g_ar_as10, N_mkt, ...
                    rand_seed, tau_reps, sig_level, ...
                    correlated_fc, data_outcome_vec2, N_dev2, corr_sim, ...
                    dev2_ind_f1_tr, dev2_ind_f2_tr, corr_g_ar);
    
                Tnmesh_fy_as10(n) = Tn; cvmesh_fy_as10(n) = cv;
            end
        
            if correlated_fc
                covered_mesh_fy_as10(:, :, :, m) = (Tnmesh_fy_as10<cvmesh_fy_as10);
            else            
                covered_mesh_fy_as10(:, :, m) = (Tnmesh_fy_as10<cvmesh_fy_as10);
            end
        end
    
        profit_lb_vec = []; profit_ub_vec = [];  data_outcome_vec = []; 
    
        %% CT and AT estimation
        fc_sim = randn(1, N_mkt, N, N_sim);    
    
        g_ar_ct_emp = {gen_iv_monte_carlo(profit_ub, profit_lb, profit_cutoff_min, profit_cutoff_max)};
        
        % (1) At the true parameter value
        % CT bds, AS test statistic
        tic;
        [Tn00ct, cv00ct] = obj_ct_as10(para00, profit, ...
            firm_outcome, ind_unilateral_deivation, fc_sim, is_equi_data, N, ....
            g_ar_ct_emp, rand_seed, tau_reps, sig_level, ...
            correlated_fc, corr_sim_ct);
    
        if monte_carlo_exp == 1
            is_covered_ct_as10(m) = Tn00ct<cv00ct;
        else
            ct_t_as10(nn, m) = toc;
        end

    
        % AT bds, AS test
        tic;
        [Tn00at, cv00at] = obj_at_as10(para00, profit_lb, profit_ub, ... 
            firm_outcome, is_equi_data, N_mkt, ....
            g_ar_ct_emp, rand_seed, tau_reps, sig_level, ...
            correlated_fc, corr_sim_at);
        
        if monte_carlo_exp == 1
            is_covered_at_as10(m) = Tn00at<cv00at;
        else
            at_t_as10(nn, m) = toc;
        end
        
        % (2) at the parameter grid points
        if monte_carlo_exp==1
            Tnmeshct = nan(size(Xmesh)); cvmeshct = nan(size(Xmesh));
            Tnmeshat = nan(size(Xmesh)); cvmeshat = nan(size(Xmesh));
    
            parfor n = 1 : size(para_grid, 1)
                [Tn, cv] = obj_ct_as10(para_grid(n,:), profit, ...
                    firm_outcome, ind_unilateral_deivation, fc_sim, is_equi_data, N, ....
                    g_ar_ct_emp, rand_seed, tau_reps, sig_level,...
                    correlated_fc, corr_sim_ct);
    
                Tnmeshct(n) = Tn; cvmeshct(n) = cv;
    
                [Tn, cv] = obj_at_as10(para_grid(n,:), profit_lb, profit_ub, ...
                    firm_outcome, is_equi_data, N_mkt, ....
                    g_ar_ct_emp, rand_seed, tau_reps, sig_level, ...
                    correlated_fc, corr_sim_at);
    
                Tnmeshat(n) = Tn; cvmeshat(n) = cv;            
            end
    
            if correlated_fc
                covered_mesh_ct_as10(:, :, :, m) = (Tnmeshct<cvmeshct);
                covered_mesh_at_as10(:, :, :, m) = (Tnmeshat<cvmeshat);
            else
                covered_mesh_ct_as10(:, :, m) = (Tnmeshct<cvmeshct);
                covered_mesh_at_as10(:, :, m) = (Tnmeshat<cvmeshat);
            end
        end
    end
end

if monte_carlo_exp == 1
    % save results, which will be loaded by plot_CoverageProb.m to generate Fig 1 and Fig 2
    tmp = sprintf('MonteCarloEstResult_N%1.0f_Nmonte%1.0f_Nmkt%1.0f_phi%1.2f.mat', N, N_monte, N_mkt, phi);
    save(tmp, ...
        'Xgrid','Ygrid','Zgrid','Xmesh','Ymesh','Zmesh','para00',...
        'is_covered_fy_as10', 'is_covered_ct_as10', 'is_covered_at_as10',...
        'covered_mesh_fy_as10', 'covered_mesh_ct_as10', 'covered_mesh_at_as10');
else
    % Table 1
    fileID_time_table = fopen("time_used.txt", 'w+');
    fprintf(fileID_time_table, 'N\tFY\tCT\tAT\n');
    for nn = 1:length(N_vec)
        fprintf(1, '%1.0f\t%1.3f\t%1.3f\t%1.3f\n', [N_vec(nn), mean(fy_t_as10(nn, 2:end),2),mean(ct_t_as10(nn,2:end)), mean(at_t_as10(nn,2:end))]);
    end
    fclose(fileID_time_table)
end